'''
 ==================板块功能描述====================
           @Time     :2022/4/9   15:34
           @Author   : qiaofengsheng
           @File     :make_mask_data.py
           @Software :PyCharm
           @description:
 ================================================
 '''
import os

import cv2
import numpy as np
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import json

CLASS_NAMES = ['horse', 'person']


def make_mask(image_dir, save_dir):
    data = os.listdir(image_dir)
    temp_data = []
    for i in data:
        if i.split('.')[1] == 'json':
            temp_data.append(i)
        else:
            continue
    for js in temp_data:
        json_data = json.load(open(os.path.join(image_dir, js), 'r'))
        shapes_ = json_data['shapes']
        mask = Image.new('P', Image.open(os.path.join(image_dir, js.replace('json', 'png'))).size)
        for shape_ in shapes_:
            label = shape_['label']
            points = shape_['points']
            points = tuple(tuple(i) for i in points)
            mask_draw = ImageDraw.Draw(mask)
            mask_draw.polygon(points, fill=CLASS_NAMES.index(label) + 1)
        mask.save(os.path.join(save_dir, js.replace('json', 'png')))


def vis_label(img):
    img=Image.open(img)
    # img=np.array(img)
    img_array = np.array(img)

    # 检查并确保转换后的数组适合matplotlib显示（如果是RGB或RGBA格式）
    if img.mode == 'RGB':
        # RGB图像，3通道
        pass
    elif img.mode == 'P':
        img = img.convert('RGB')
    elif img.mode == 'RGBA':
        # RGBA图像，4通道，需要将alpha通道考虑进去
        img_array = img_array[..., :3]  # 如果不需要alpha通道，可以去掉它
    else:
        # 灰度图或其他单通道图像，可能需要转换为RGB
        img_array = img_array.reshape(img.size[1], img.size[0], 1)  # 转换为单通道三维数组
        img_array = np.concatenate((img_array, img_array, img_array), axis=2)  # 转换为RGB图像
    # 使用matplotlib显示图像
    plt.figure(figsize=(8, 6))
    plt.imshow(img_array)
    plt.axis('off')  # 可选，移除坐标轴
    plt.show()
    print(img)
    # print(set(img.reshape(-1).tolist()))



if __name__ == '__main__':
    # make_mask('image', 'SegmentationClass')
    vis_label('SegmentationClass/2007_000032.png')
    # img=Image.open('SegmentationClass/000019.png')
    # print(np.array(img).shape)
    # out=np.array(img).reshape(-1)
    # print(set(out.tolist()))
